﻿// InsertSection.cpp: 实现文件
//

#include "pch.h"
#include "MY_PE_TOOL.h"
#include "InsertSection.h"
#include "afxdialogex.h"
#include "MyHeadFile.h"


// InsertSection 对话框

IMPLEMENT_DYNAMIC(InsertSection, CDialogEx)

InsertSection::InsertSection(CWnd* pParent /*=nullptr*/)
	: CDialogEx(IDD_DIALOG4, pParent)
{

}

InsertSection::~InsertSection()
{
}

void InsertSection::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Control(pDX, EDT_SIZE, m_EditSecSize);
}


DWORD InsertSection::CalcOffset(DWORD dwVal, DWORD dwAlign)
{
    if (dwVal % dwAlign == 0)
    {
        return dwVal;
    }
    else
    {
        return (dwVal / dwAlign + 1) * dwAlign;
    }

    return 0;
}

//扩充节
void InsertSection::IncSection(DWORD dwValue)
{
	DWORD dwWriteBytes = 0;

    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_SECTION_HEADER pSectionHead = (PIMAGE_SECTION_HEADER)((DWORD)&pNtHead->OptionalHeader + pNtHead->FileHeader.SizeOfOptionalHeader);

	DWORD dwFileAli = CalcOffset(dwValue, pNtHead->OptionalHeader.FileAlignment);

	//最后一个节的下标索引
	DWORD dwSectionIndex = pNtHead->FileHeader.NumberOfSections - 1;

	DWORD dwOldMemAli = CalcOffset(pSectionHead[dwSectionIndex].Misc.VirtualSize, pNtHead->OptionalHeader.SectionAlignment);

	IMAGE_SECTION_HEADER SecHead = { 0 };
	memcpy(&SecHead, &pSectionHead[dwSectionIndex], sizeof(IMAGE_SECTION_HEADER));

	SecHead.SizeOfRawData = pSectionHead[dwSectionIndex].SizeOfRawData + dwFileAli;
	SecHead.Misc.VirtualSize = pSectionHead[dwSectionIndex].Misc.VirtualSize + dwFileAli;

	DWORD dwNewMemAli = CalcOffset(SecHead.Misc.VirtualSize, pNtHead->OptionalHeader.SectionAlignment);
    //打开文件
    HANDLE hFile = CreateFile(m_pFileFullPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (hFile == INVALID_HANDLE_VALUE)
    {
        MessageBox(_T("打开文件失败"), _T("提示"), MB_OK);
        return;
    }

	//计算是否需要修改SizeofImage
	if (dwNewMemAli > dwOldMemAli)
	{
        //计算SizeofImage偏移
		DWORD dwost = (DWORD)&pNtHead->OptionalHeader.SizeOfImage - (DWORD)pDosHead;
        ::SetFilePointer(hFile, dwost, NULL, FILE_BEGIN);

        //写入节表项
		dwost = pNtHead->OptionalHeader.SizeOfImage + dwNewMemAli - dwOldMemAli;
        ::WriteFile(hFile, &dwost, 4, &dwWriteBytes, NULL);
	}

    //计算节的偏移
	DWORD dwost = (DWORD)(&pSectionHead[dwSectionIndex].Misc.VirtualSize) - (DWORD)pDosHead;
    ::SetFilePointer(hFile, dwost, NULL, FILE_BEGIN);

    //写入节表项
    ::WriteFile(hFile, &SecHead.Misc.VirtualSize, 4, &dwWriteBytes, NULL);

    dwost = (DWORD)(&pSectionHead[dwSectionIndex].SizeOfRawData) - (DWORD)pDosHead;
    ::SetFilePointer(hFile, dwost, NULL, FILE_BEGIN);

    //写入节表项
    ::WriteFile(hFile, &SecHead.SizeOfRawData, 4, &dwWriteBytes, NULL);

    //写入节数据
    ::SetFilePointer(hFile, 0, NULL, FILE_END);
    char* pBuff = new char[dwFileAli];
    ZeroMemory(pBuff, dwFileAli);
    WriteFile(hFile, pBuff, dwFileAli, &dwWriteBytes, NULL);

    delete[] pBuff;
    //关闭文件
    CloseHandle(hFile);
}

//增加节
void InsertSection::AddSection(DWORD dwValue)
{
	DWORD dwWriteBytes = 0;

    PIMAGE_DOS_HEADER pDosHead = (PIMAGE_DOS_HEADER)(*m_pFileData);
    PIMAGE_NT_HEADERS pNtHead = (PIMAGE_NT_HEADERS)((DWORD)pDosHead + pDosHead->e_lfanew);
    PIMAGE_SECTION_HEADER pSectionHead = (PIMAGE_SECTION_HEADER)((DWORD)&pNtHead->OptionalHeader + pNtHead->FileHeader.SizeOfOptionalHeader);
	
	DWORD dwFileAli = CalcOffset(dwValue, pNtHead->OptionalHeader.FileAlignment);

	DWORD dwSectionIndex = pNtHead->FileHeader.NumberOfSections - 1;

	IMAGE_SECTION_HEADER SecHead = { 0 };

	SecHead.Misc.VirtualSize = dwFileAli;
	DWORD dwMemAli = CalcOffset(pSectionHead[dwSectionIndex].Misc.VirtualSize, pNtHead->OptionalHeader.SectionAlignment);
	SecHead.VirtualAddress = pSectionHead[dwSectionIndex].VirtualAddress + dwMemAli;
	SecHead.SizeOfRawData = dwFileAli;
	SecHead.PointerToRawData = pSectionHead[dwSectionIndex].PointerToRawData + pSectionHead[dwSectionIndex].SizeOfRawData;
	SecHead.Characteristics = IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE;

	//打开文件
    HANDLE hFile = CreateFile(m_pFileFullPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (hFile == INVALID_HANDLE_VALUE)
    {
        MessageBox(_T("打开文件失败"), _T("提示"), MB_OK);
        return;
    }


    //计算节表偏移
	DWORD dwost = (DWORD)&pNtHead->FileHeader.NumberOfSections - (DWORD)pDosHead;
    ::SetFilePointer(hFile, dwost, NULL, FILE_BEGIN);

	//写入节表项
	DWORD dwNewSec = pNtHead->FileHeader.NumberOfSections + 1;
	::WriteFile(hFile, &dwNewSec, 4, &dwWriteBytes, NULL);

	//计算SizeofImage偏移
	dwost = (DWORD)&pNtHead->OptionalHeader.SizeOfImage - (DWORD)pDosHead;
	::SetFilePointer(hFile, dwost, NULL, FILE_BEGIN);

	//写入节表项
	dwost = CalcOffset(dwValue, pNtHead->OptionalHeader.SectionAlignment) + pNtHead->OptionalHeader.SizeOfImage;
	::WriteFile(hFile, &dwost, 4, &dwWriteBytes, NULL);

	//将节区数据写入节表项
	::SetFilePointer(hFile, m_Offset, NULL, FILE_BEGIN);
    BOOL bRet = WriteFile(hFile, SecHead.Name, 8, &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.Misc.VirtualSize, sizeof(SecHead.Misc.VirtualSize), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.VirtualAddress, sizeof(SecHead.VirtualAddress), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.SizeOfRawData, sizeof(SecHead.SizeOfRawData), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.PointerToRawData, sizeof(SecHead.PointerToRawData), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.PointerToRelocations, sizeof(SecHead.PointerToRelocations), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.PointerToLinenumbers, sizeof(SecHead.PointerToLinenumbers), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.NumberOfRelocations, sizeof(SecHead.NumberOfRelocations), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.NumberOfLinenumbers, sizeof(SecHead.NumberOfLinenumbers), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.Characteristics, sizeof(SecHead.Characteristics), &dwWriteBytes, NULL);


	//写入节数据
	::SetFilePointer(hFile, 0, NULL, FILE_END);
	char *pBuff = new char[dwFileAli];
	ZeroMemory(pBuff, dwFileAli);
	WriteFile(hFile, pBuff, dwFileAli, &dwWriteBytes, NULL);

	delete[] pBuff;
	//关闭文件
	CloseHandle(hFile);
}

//写文件
void InsertSection::MyWriteFile(IMAGE_SECTION_HEADER SecHead)
{
    //打开文件
    HANDLE hFile = CreateFile(m_pFileFullPath, GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

    if (hFile == INVALID_HANDLE_VALUE)
    {
        MessageBox(_T("打开文件失败"), _T("提示"), MB_OK);
        return;
    }

    //设置文件指针偏移
    DWORD D = SetFilePointer(hFile, m_Offset, NULL, FILE_BEGIN);

    //写文件
	DWORD dwWriteBytes = 0;
    BOOL bRet = WriteFile(hFile, SecHead.Name, 8, &dwWriteBytes, NULL);
	bRet = WriteFile(hFile, &SecHead.Misc.VirtualSize, sizeof(SecHead.Misc.VirtualSize), &dwWriteBytes, NULL);
	bRet = WriteFile(hFile, &SecHead.VirtualAddress, sizeof(SecHead.VirtualAddress), &dwWriteBytes, NULL);
	bRet = WriteFile(hFile, &SecHead.SizeOfRawData, sizeof(SecHead.SizeOfRawData), &dwWriteBytes, NULL);
	bRet = WriteFile(hFile, &SecHead.PointerToRawData, sizeof(SecHead.PointerToRawData), &dwWriteBytes, NULL);
	bRet = WriteFile(hFile, &SecHead.PointerToRelocations, sizeof(SecHead.PointerToRelocations), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.PointerToLinenumbers, sizeof(SecHead.PointerToLinenumbers), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.NumberOfRelocations, sizeof(SecHead.NumberOfRelocations), &dwWriteBytes, NULL);
    bRet = WriteFile(hFile, &SecHead.NumberOfLinenumbers, sizeof(SecHead.NumberOfLinenumbers), &dwWriteBytes, NULL);
	bRet = WriteFile(hFile, &SecHead.Characteristics, sizeof(SecHead.Characteristics), &dwWriteBytes, NULL);

    //关闭句柄
    CloseHandle(hFile);
    return;
}


void InsertSection::SetSectionType(DWORD SecType)
{
	m_SectionType = SecType;
}

void InsertSection::SetFilePath(TCHAR* pFilePath)
{
	m_pFileFullPath = pFilePath;
}

void InsertSection::SetFileData(char** pFileData)
{
	m_pFileData = pFileData;
}

void InsertSection::SetOffset(DWORD dwOffset)
{
	m_Offset = dwOffset;
}

BEGIN_MESSAGE_MAP(InsertSection, CDialogEx)
	ON_BN_CLICKED(BTN_OK, &InsertSection::OnBnClickedOk)
END_MESSAGE_MAP()


// InsertSection 消息处理程序


//点击确定按钮
void InsertSection::OnBnClickedOk()
{
	CString strSize;
	m_EditSecSize.GetWindowTextW(strSize);

	if (strSize.IsEmpty())
	{
		MessageBox(_T("增加失败"), _T("提示"), MB_OK);
		return;
	}

    //将输入的值转为数字
    DWORD dwValue = wcstoul(strSize.GetBuffer(), NULL, 16);

	//扩充节
	if (m_SectionType == INC_SECTION)
	{
		IncSection(dwValue);
	}
	//增加节
	else
	{
		AddSection(dwValue);
	}

	PostMessageW(WM_CLOSE);
}


BOOL InsertSection::OnInitDialog()
{
	CDialogEx::OnInitDialog();

	if (m_SectionType == INC_SECTION)
	{
		SetDlgItemTextW(STC_MYSELF, _T("扩充节的大小："));
	}
	else
	{
		SetDlgItemTextW(STC_MYSELF, _T("增加节的大小："));
	}

	return TRUE;  // return TRUE unless you set the focus to a control
				  // 异常: OCX 属性页应返回 FALSE
}
